import numpy as np
from typing import List, Tuple, Union

class BBox:
    def __init__(self, x1, y1, x2, y2):
        """
        使用 xyxy 格式初始化边界框。
        :param x1: float，左上角 x 坐标
        :param y1: float，左上角 y 坐标
        :param x2: float，右下角 x 坐标
        :param y2: float，右下角 y 坐标
        """
        self.x1 = x1
        self.y1 = y1
        self.x2 = x2
        self.y2 = y2

    @classmethod
    def from_xywh(cls, x, y, w, h):
        """
        通过 xywh 格式创建边界框。
        :param x: float，左上角 x
        :param y: float，左上角 y
        :param w: float，宽度
        :param h: float，高度
        :return: BBox 实例
        """
        return cls(x, y, x + w, y + h)

    @property
    def xyxy(self):
        """返回 xyxy 格式：[x1, y1, x2, y2]"""
        return [self.x1, self.y1, self.x2, self.y2]

    @property
    def xywh(self):
        """返回 xywh 格式：[x, y, w, h]"""
        return [self.x1, self.y1, self.x2 - self.x1, self.y2 - self.y1]

    def __mul__(self, scale):
        """
        缩放边界框坐标。
        :param scale: tuple (width, height)
        :return: 缩放后的 BBox
        """
        w, h = scale
        return BBox(self.x1 * w, self.y1 * h, self.x2 * w, self.y2 * h)
        
    def __truediv__(self, scale):
        """
        缩小边界框坐标（除法）。
        :param scale: tuple (width, height)
        :return: 缩放后的 BBox（归一化）
        """
        w, h = scale
        return BBox(self.x1 / w, self.y1 / h, self.x2 / w, self.y2 / h)

    def __repr__(self):
        return f"BBox(xyxy={self.xyxy}, xywh={self.xywh})"

    def to_yolo_format(self, image_size):
        """
        转换为 YOLO 格式 (cx, cy, w, h)，坐标归一化。
        :param image_size: tuple (image_width, image_height)
        :return: list [cx, cy, w, h]，归一化后的 YOLO 格式
        """
        img_w, img_h = image_size
        x1, y1, x2, y2 = self.xyxy
        cx = (x1 + x2) / 2 / img_w
        cy = (y1 + y2) / 2 / img_h
        w = (x2 - x1) / img_w
        h = (y2 - y1) / img_h
        return [cx, cy, w, h]

class PolyMask:
    def __init__(self, points: Union[List[float], List[Tuple[float, float]]]):
        """
        初始化多边形掩码。
        :param points: 一维 list [x1, y1, x2, y2, ...] 或 list of (x, y) tuples
        """
        if isinstance(points, list) and all(isinstance(p, (int, float)) for p in points):
            if len(points) % 2 != 0:
                raise ValueError("一维 points 列表长度必须为偶数")
            self._points = [(points[i], points[i + 1]) for i in range(0, len(points), 2)]
        elif isinstance(points, list) and all(isinstance(p, (list, tuple)) and len(p) == 2 for p in points):
            self._points = [tuple(p) for p in points]
        else:
            raise TypeError("points 应为一维 list 或 list of (x, y) tuples")

    def __mul__(self, scale: Tuple[float, float]) -> 'PolyMask':
        """
        缩放多边形坐标。
        :param scale: (width, height)
        :return: 缩放后的 PolyMask 实例
        """
        if not isinstance(scale, (list, tuple)) or len(scale) != 2:
            raise ValueError("scale 应为 (width, height) 的二元组")
        w, h = scale
        scaled_points = [(x * w, y * h) for x, y in self._points]
        return PolyMask(scaled_points)

    def __truediv__(self, scale: Tuple[float, float]) -> 'PolyMask':
        """
        归一化多边形坐标（除法）。
        :param scale: (width, height)
        :return: 归一化后的 PolyMask
        """
        if not isinstance(scale, (list, tuple)) or len(scale) != 2:
            raise ValueError("scale 应为 (width, height) 的二元组")
        w, h = scale
        scaled_points = [(x / w, y / h) for x, y in self._points]
        return PolyMask(scaled_points)

    def __len__(self) -> int:
        return len(self._points)

    def __repr__(self) -> str:
        return f"PolyMask(num_points={len(self._points)})"

    def bbox(self) -> 'BBox':
        """
        返回包含该多边形的最小边界框。
        """
        points = np.array(self._points)
        x_min = float(np.min(points[:, 0]))
        y_min = float(np.min(points[:, 1]))
        x_max = float(np.max(points[:, 0]))
        y_max = float(np.max(points[:, 1]))
        return BBox(x_min, y_min, x_max, y_max)

    @property
    def points(self) -> List[Tuple[float, float]]:
        """
        返回原始点坐标列表。
        """
        return self._points

    def flat(self) -> List[float]:
        """
        返回平铺的一维坐标列表 [x1, y1, x2, y2, ...]
        """
        return [coord for pt in self._points for coord in pt]
        
    def normalize_in_bbox(self, bbox: 'BBox') -> 'PolyMask':
        """
        将多边形归一化到指定的 BBox 中（相对坐标，范围 [0, 1]）。
        :param bbox: BBox 对象
        :return: 归一化后的 PolyMask
        """
        x1, y1, x2, y2 = bbox.xyxy
        if x2 - x1 == 0 or y2 - y1 == 0:
            raise ValueError("无效的 BBox，宽或高为 0")
    
        normalized_points = [((x - x1) / (x2 - x1), (y - y1) / (y2 - y1)) for x, y in self._points]
        return PolyMask(normalized_points)

    def clip_to_bbox(self, bbox: 'BBox') -> 'PolyMask':
        """
        将多边形的每个点限制在指定的 BBox 内。
        :param bbox: BBox 对象
        :return: 裁剪后的 PolyMask
        """
        x1, y1, x2, y2 = bbox.xyxy
        clipped_points = [
            (min(max(x, x1), x2), min(max(y, y1), y2))
            for x, y in self._points
        ]
        return PolyMask(clipped_points)
    
class Anomaly:
    def __init__(self, score, bbox, polygon, category):
        """
        缺陷标注类。
        :param score: float，置信度
        :param bbox: BBox 实例
        :param polygon: PolyMask 实例或 None
        :param category: int 或 str，类别信息
        """
        self.score = score
        if isinstance(bbox, list):
            self.bbox = BBox(*bbox)
        else:
            self.bbox = bbox
        if isinstance(polygon, list):
            self.polygon = PolyMask(polygon)
        else:
            self.polygon = polygon

        if self.polygon is not None:
            self.polygon = self.polygon.clip_to_bbox(self.bbox)

        self.category = category

    def __repr__(self):
        return (
            f"Anomaly(score={self.score:.2f}, "
            f"bbox={self.bbox}, "
            f"polygon_len={len(self.polygon) if self.polygon else 0}, "
            f"category={self.category})"
        )


def parse_labels(ann_file, with_conf=False, task="segment"):
    results = []
    with open(ann_file) as f:
        annotations = f.readlines()

    for item in annotations:
        item = item.split(" ")
        class_id = item[0]

        if task == "segment":
            if with_conf:
                if len(item[1:]) < 2:
                    continue
                coordinates = list(map(float, item[1:-1]))
                score = float(item[-1])
            else:
                coordinates = list(map(float, item[1:]))
                score = 1
            polygon_points = PolyMask(coordinates)
            bbox = polygon_points.bbox()
        elif task == "detect":
            if with_conf:
                x,y,w,h = list(map(float, item[1:-1]))
                x1, y1, x2, y2 = [x-w/2, y-h/2, x+w/2, y+h/2]
                score = float(item[-1])
                bbox = BBox(x1, y1, x2, y2)
                polygon_points = None
            else:
                x,y,w,h = list(map(float, item[1:]))
                x1, y1, x2, y2 = [(x-w/2, y-h/2), (x+w/2, y-h/2), (x+w/2, y+h/2), (x-w/2, y+h/2)]
                score = 1
                bbox = BBox(x1, y1, x2, y2)
                polygon_points = None


        flaw = Anomaly(score, bbox, polygon_points, class_id)
        results.append(flaw)
    return results

def yolo_to_coco(yolo_dir, save_path, split=''):
    images_dir = osp.join(yolo_dir, 'images', split)
    labels_dir = osp.join(yolo_dir, 'labels', split)
    classes_path = osp.join(yolo_dir, 'classes.txt')
    with open(classes_path, 'r') as f:
        class_names = [line.strip() for line in f.readlines()]
    categories = [{"id": i, "name": name} for i, name in enumerate(class_names)]
    images = []
    annotations = []
    ann_id = 0
    image_id = 0
    for img_name in tqdm(os.listdir(images_dir)):
        if not img_name.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            continue
        image_id += 1
        img_path = os.path.join(images_dir, img_name)
        img_arr = cv2.imread(img_path)
        h, w = img_arr.shape[:2]
        images.append({
            "id": image_id,
            "file_name": img_name,
            "width": w,
            "height": h
        })

        label_path = os.path.join(labels_dir, osp.splitext(img_name)[0] + '.txt')
        labels = parse_labels(label_path)
        for item in labels:
            ann_id += 1
            annotations.append({
                "id": ann_id,
                "image_id": image_id,
                "category_id": int(cls_id),
                "bbox": item.bbox.xywh * (w, h),
                "segmentation": [item.polygon * (w, h)],
                "iscrowd": 0
            })
    coco_dict = {
        "images": images,
        "annotations": annotations,
        "categories": categories
    }
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    with open(save_path, 'w') as f:
        json.dump(coco_dict, f, indent=2)
    print(f"✅ COCO JSON saved to: {save_path}")